#!/usr/bin/env python

from ssummo.cmd_options import Options
from ssummo import colours
from ssummo.graphs import Plotter, make_patch_spines_invisible, \
                          make_spine_invisible
from ssummo.resample import Indices, RandomNumberThread, Replicate, Results
from ssummo.ssummolib import collapse_at_rank, load_index
from ssummo.taxonomy import TaxDB

import re
import multiprocessing
import os.path
import Queue
import sys

try:
    import cPickle as pickle
except ImportError:
    import pickle

try:
    import matplotlib.pyplot as plt
    matplot = True
except ImportError:
    matplot = False
    sys.stderr.write("Matplotlib import error....")
    sys.stderr.write("  Graph-plotting functionality will not work. :'(")

class LocalOptions(Options):
    def __init__(self, args=None):
        super(LocalOptions, self).__init__()
        self.singleargs = ['--draw', '-replicates',
                           '-increment', '-collapse-at-rank']
        self.multiargs = ['-in', '-out', '-order', '-indices']
        self.help_text = { '-in' : 'Input files.',
                         '-out': 'Output files. Default: standard out',
                         '-replicates': 'Number of replicates for each dataset. [5]',
                         '-increment' : 'For rarefaction curves, only sample every `-increment` species. [100]',
                         '--draw' : 'If results are already saved, just draw the results saved in file `-in`',
                         '-collapse-at-rank': 'Collapse results at specific rank.',
                         '--indices'  :  'Calculate and plot ALL biodiversity indices.',
                         '-indices'   :  ' List which indices to plot [shannon | shannon_max | simpson | jack].',
                         '-order' : 'Defines the order when drawing old results.',
                         }
        self.options.update( { '-replicates' : None    ,
                               '-increment'  : None    ,
                               '--draw'      : False   ,
                               '-collapse-at-rank': None,
                               '--indices'   : False   ,
                               '-indices'    : []      ,
                               '-out' : []             ,
                             } )
        self.groups = None
        self.switches = {'--draw' : False ,
                         '--indices' : False}
        self.post_checks += [ self.groupify , self.indices]
        if args is not None:
            self.parse_args(args)

    def groupify(self):
        """If there are any groups defined on the command line, this'll just
        create a list contianing the lengths of each group.
        """
        if len( self.regopts ) > 0:
            self.groups = [len(in_files) for in_files in \
                          [self.regopts[key] for key in \
                            sorted(self.regopts.keys())]]

    def indices(self):
        """If you specify --indices, append all."""
        if self.options['--indices']:
            self.options['-indices'] = ['shannon', 'shannon_max',
                                        'simpson', 'jack', ]
        elif len(self.options['-indices']) > 0:
            self.options['--indices'] = True

class PlotRarefaction(Plotter):
    def __init__(self, options):
        # datasets, replicates=None, increment=None, groups=None):
        """
        """
        if options['-replicates'] is None:
            self.replicates = 5
        else:
            self.replicates = int(options['-replicates'])
        if options['-increment'] is None:
            self.increment = 1000
        else:
            self.increment = int(options['-increment'])

        self.col_ind = 0   # For assigning colours.
        number_datasets =  len(options['-in'])
        self.n = number_datasets
        self.dataset_names = []
        self.options = options
        if '_groups' in self.options.options:
            self.colours = [colours.HSL_to_HEX(colour) for colour in 
                            colours.Heat(N=self.n,
                                         groups=self.options['_groups'])]
        else:
            self.colours = [colours.HSL_to_HEX(colour) for colour in
                            colours.Heat(N=self.n)]

        if matplot:
            self.fig = plt.figure()
            self.axis = self.fig.add_subplot( 111 )
            plt.ioff()
            # Turn off interactive mode (default anyway, unless changed in
            # matplotlib config file); speeds things up.
            plt.xlabel( 'Number of randomly chosen sequences' )
            plt.ylabel( 'Number of distinct taxa' )
            if self.options['--indices']:
                self.init_axes()

    def reinit(self):
        """To be called if need to reload the colours. This needs doing with
        --draw, when PlotRarefaction is initiated before the previous results
        have been loaded (i.e. before the original options have been loaded).
        """
        self.n = len(self.options['-in'])
        print(self.options['_groups'])
        self.colours = [colours.HSL_to_HEX(colour) for colour in \
                        colours.Heat(N=self.n, groups=self.options['_groups'])]

    def init_axes(self):
        """Here define methods to initialise additional axes. This is called
        during initialisation if options['--indices'] is True.
        """
        axis_adjustment = 0.05
        setattr( self , 'fig2' , plt.figure() )
        setattr( self, 'axis2' , self.fig2.add_subplot( 111 ) )
        self.axis2.set_xlabel( 'Number sequences sampled' )
        cur_axis = self.axis2
        left_increment = 0.
        right_increment = 1.
        axis_ind = 3
        n_axes = len( self.options['-indices'] )
        if n_axes > 1:
            side = 'right'
            for counter in xrange( n_axes ) :
                # Move the y-axes across.
                if side == 'left':
                    new_axis = self.fig2.add_subplot(111, axisbg='#000000')
                    left_increment -= axis_adjustment
                    new_axis.spines[side].set_position(('axes', left_increment))
                    make_patch_spines_invisible(new_axis)
                    make_spine_invisible(new_axis, side)
                    side = 'right'
                elif side == 'right':
                    new_axis = cur_axis.twinx()
                    #if axis_ind > 3:
                    new_axis.spines[side].set_position(('axes', right_increment))
                    make_patch_spines_invisible(new_axis)
                    make_spine_invisible(new_axis, side)
                    #else:
                    #    new_axis.spines['left'].set_position( ('axes',right_increment ) )
                    #make_patch_spines_invisible( new_axis )
                    #    make_spine_invisible( new_axis ,'left')
                    right_increment += axis_adjustment
                    side = 'left'
                else:
                    raise IOError

                axis_id = 'axis{0}'.format( axis_ind )
                print('creating ', axis_id)
                setattr( self , axis_id , new_axis )
                axis_ind += 1

                cur_axis = new_axis

        self.fig2.subplots_adjust(left=0.15, right=0.85)
        self.axes = range(2, axis_ind)
        self.cur_axis = 2

    def draw(self):
        """Just provide the x-axis and y-axis limits. The x-axis is rounded
        up to the nearest 100, and the y-axis is rounded up to the nearest
        500.
        """
        self.axis.legend(self.artists, self.options['-in'],
                         loc=4, borderaxespad=0.2)
        if self.options['--indices']:
            self.draw_indices(self.index_artists, self.index_labels)

        # round up to the nearest 500
        self.axis.set_xlim( 0, 500 * (1 + (self.x_max // 500)) )
        # round up to the nearest 100
        self.axis.set_ylim( 0, 100 * (1 + (self.y_max // 100)) )

        ticker = plt.matplotlib.ticker.MultipleLocator(base=self.increment * 2)
        subticker = plt.matplotlib.ticker.MultipleLocator(base=self.increment)
        (self.ticker, self.subticker, ) = (ticker, subticker, )

        self.axis.xaxis.set_major_locator(ticker)
        self.axis.xaxis.set_minor_locator(subticker)
        if self.options['--indices']:
            i = 2
            while hasattr(self, 'axis{0}'.format(i)):
                ax = getattr(self, 'axis{0}'.format(i))
                ax.xaxis.set_major_locator(ticker)
                ax.xaxis.set_minor_locator(subticker)
                i += 1
        self.axis.grid(True)
        plt.show()

    def draw_indices(self, index_artists, index_labels):
        self.axis2.legend( index_artists, index_labels, loc=4, borderaxespad=0.2 )
        index = 2
        for i, bio_index in enumerate( self.options['-indices'] ):
            axis_id = 'axis{0}'.format( index )
            if hasattr( self , axis_id ):
                print('setting {0} on {1}'.format(bio_index, axis_id))
                cur_axis = getattr( self , axis_id )
                cur_axis.set_ylabel( bio_index )
                cur_axis.tick_params(axis='y', size=4, width=1.5 )
                index += 1
            else:
                print('{0} does not exist'.format(axis_id))

    def load_rarefaction(self, data):
        # Give matplotlib the data to plot behind the scenes.
        # We keep the 'artists' to colour later.
        #xvals = []; yvals = []
        #for xval in sorted( data.rarefaction.keys() ):
        #    xvals.append( xval ) ; yvals.append( data.rarefaction[xval] )

        #labels = self.options['-in']
        if len( data.rarefaction ) == 0:
            return
        if self.options['--indices']:
            self.cur_axis = 2
            sys.stdout.write( '{0}\t'.format( data.file_name ) )
            # -indices = [ 'jack' , 'shannon' , 'shannon_max' , 'simpson' ]
            for bio_index in self.options['-indices']:
                if hasattr( data , bio_index ):
                    index_data = [ getattr( data , bio_index )[xval] for xval in data.xvals ]
                    #plotter.cur_axis = plotter.axes[ plotter._marker_ind ]
                    collection = self.plot_scatter_errors( data.xvals, index_data, label=bio_index )
                    self._marker_ind += 1
                    self.index_artists.append( collection )
                    self.index_labels.append( '{0} {1}'.format( data.file_name, bio_index ) )
                    self.cur_axis += 1
            self._marker_ind = 0

        #self.x_max = max( [data.xvals , self.x_max ])
        self.x_max = max(data.xvals[-1], self.x_max)
        yvals = [ data.rarefaction[xval] for xval in data.xvals ]
        self.y_max = max(yvals[-1] + [self.y_max])
        mappings = self.plot_box(data.xvals, yvals)
        self.artists.append( mappings['boxes'][0] )

    def resume(self, results, loc=4):
        """
        :param results: Should be a dictionary of :class:`.Results` instances.
                        If you want them presented in a specific order, then
                        give the keys of the results dictionary in the order
                        which you want them displayed.
        :param loc: The location of the legend according to matplotlib
                    co-ordinates.
        """
        #artists = []
        #data_names = []
        #max_x = 0
        #max_y = 0
        order = [truncate_name(infile) for infile in self.options['-in']]
        try:
            iterator = [ results[key] for key in order ]
        except KeyError:
            print(results.keys())
            print(key)
            raise
        self.reinit()

        if self.options['-increment'] is None:
            xvals = []
            for res in iterator:
                if len(res.xvals) > 0:
                    xvals.append(max(res.xvals))
            increment = calc_increment(xvals)
            self.increment = increment

        for results_obj in iterator:
            self.load_rarefaction( results_obj )
#            print(data_file, len( results_obj.xvals ) , len(results_obj.rarefaction ))
#            xvals, yvals = results_obj.xvals , results_obj.rarefaction
#            max_x = max( max(xvals) , max_x )
#            max_y = max( max(yvals) , max_y )
#            data_names.append( data_file )
#            mappings = self.plot_box( xvals , yvals )
#            artists.append( mappings['boxes'][0] )
#        plt.legend( artists , data_names , loc=loc , borderaxespad=0.2 )
        self.draw( )
        pass

class Loader(multiprocessing.Process):
    """We have one Loader instance for each SSuMMo results file. When
    :meth:`.run` is called, (usually with Loader().start() to run as a separate
    process), a :class:`.Results` instance will be created and returned via
    :data:`.queue`, if it exists, or just returned otherwise.
    """

    def __init__(self, results_dict, options, queue=None, semaphore=None):
                 #replicates=None, increment=None):
        """This can optionally be run as a separate process, in which case
        define a queue and a semaphore in the master process, and pass them
        to Loader when initialising. e.g. ::

            import multiprocessing
            sem = multiprocessing.Semaphore()
            number_of_procs = multiprocessing.cpu_count()
            for j in range(number_of_procs):
                # do some work
                instance = Loader(results, options, multiprocessing.Queue(),
                                  semaphore=sem )
                sem.acquire()
                instance.start()

        Look at :func:`.main()` to see the full implementation.
        """
        self.tax_index = load_index()
        if queue is not None:
            multiprocessing.Process.__init__(self)
        self.queue = queue
        self.semaphore = semaphore
        self.accession_places = get_accession_places(results_dict, accessions={})
        self.nOTUs = len( set( self.accession_places.values() ) )
        self.nSeqs = len( self.accession_places )   # X-axis
        self.results_dict = results_dict
        self.options = options

        if options['-replicates'] is None:
            self.replicates = 5
        else:
            self.replicates = int(options['-replicates'])

        if options['-increment'] is None:
            self.increment = self.nSeqs / 20
        else:
            self.increment = int(options['-increment'])
        self.indices = None
        if options['--indices'] > 0:
            self.indices = True

    def _calc_indices(self, n_seqs_sampled, nodes):
        counts = sorted([len(accessions) for accessions in nodes.values()],
                        reverse=True)
        probs = [float(count) / n_seqs_sampled for count in counts]
        total = n_seqs_sampled
        indices = Indices(self.options)
        rep = Replicate()
        rep.nseqs = total
        rep.shannon, rep.shannon_max = indices.shannon_index(probs)
        rep.simpson = indices.simpsons_index(counts, total)
        rep.jack = indices.jackKnife(counts, total)
        return rep

    def run(self):
        """Loads the results in results_dict and populates a dictionary,
        with x-values as the keys, and y-values as the values.
        This allows multiple y-values to be stored per x-value.
        Optional queue will be used to send the results instead of
        returning them
        """
        results = Results()
        for number, nodes in self.rarefy():
            n_found_OTUs = len(nodes)
            if self.indices is not None:
                replicate = self._calc_indices(number, nodes)
                replicate.rarefaction = n_found_OTUs
            else:
                replicate = Replicate()
                replicate.nseqs = number
                replicate.rarefaction = n_found_OTUs
            results.append(replicate)
        if self.semaphore is not None:
            self.semaphore.release()
        if self.queue is not None:
            self.queue.put(results)
        return results

    def rarefy(self):
        """Given a results dictionary, this will sample random accessions
        from that dictionary and will _yield_ the number of accessions
        sampled as well as the number of genera / species / OTUs that a random
        sample of that number of accessions finds.

        Therefore, this is an interator generator function, yielding the
        tuple ``(nsampled, found_nodes)`` for each sample replicate.
        """
        node = self.results_dict
        if self.options['-collapse-at-rank'] is not None:
            node = collapse_at_rank(node, self.options, TaxDB(), self.tax_index)
        accessions = self.accession_places.keys()
        n = self.nSeqs
        nsampled = self.increment
        random_queue = Queue.Queue()
        random_generator = RandomNumberThread(random_queue, n, self.replicates,
                                              self.increment)
        random_generator.start()

        while nsampled < n:
            for replicate in xrange( self.replicates ):
                running_accessions = accessions[:]
                acc_places = self.accession_places.copy()
                n_left = n
                found_nodes = {}
                for ii in xrange( nsampled ):
                    accession = running_accessions.pop(random_queue.get())
                    found_node = acc_places.pop( accession )
                    if found_node in found_nodes:
                        found_nodes[found_node].append(accession)
                    else:
                        found_nodes.update({ found_node : [accession] })
                    n_left -= 1   # Every time we choose a random accession, the list of accessions is 1 shorter.
                    random_queue.task_done()
                yield nsampled, found_nodes
            nsampled += self.increment
        assert random_queue.get() == 'done'
        random_queue.task_done()
        random_generator.join()

def calc_increment(n_sequences, steps=None, max_fn=max):
    """Calculates a decent (x-axis) increment. Will by default figure out the
    increment widths for 20 steps, for a number of sequences determined by
    :func:`max_fn`.

    :param n_sequences: Should be a sequence containing the number of sequences
                        for each set of results being plotted.

    :param steps: Number of increments desired for the graph (default=20)

    :param max_fn: A callback function to be called with :data:`.n_sequences`.
                   The returned value will be divided by `steps` to calculate
                   the increment.

    e.g. ::

        >>> import numpy as np
        >>> # lambda function to calculate 95% confidence limits.
        >>> my_maximum = lambda x : np.average(x) +  (1.96 * np.std(x) )
        >>> # generate some random data
        >>> n_seqs = [np.random.random(10) for i in xrange(10)]
        >>> increment = calc_increment(n_seqs, steps=40, max_fn=my_maximum)
    """
    if steps is None:
        steps = 20
    #seqs = []
    maximum = max_fn(n_sequences)
    increment = maximum / steps
    return increment

def get_accession_places(node, accessions={}, path='/'):
    join = os.path.join
    accs = accessions
    for key in node.keys():
        if key == 'accessions':
            for acc in node[key]:
                accs.update({ acc : path })
        else:
            accs = get_accession_places(node[key], accessions=accs,
                                        path=join(path, key))
    return accs

def get_save_name():
    """Asks a user for a file name. If it exists, ask again,
    if user types nothing, ask again...
    Return None if user types 'n' or 'no'.
    """
    save_name = raw_input('Please enter a name to save the results, which can be drawn later:-\n')
    if save_name == ('n' or 'no'):
        return None
    elif os.path.exists(save_name):
        print('That already exists. Trying again. (`no` or `n` to skip saving)')
        return get_save_name()
    elif save_name == '':
        print("If you wanna leave without saving, say 'no'!")
        return get_save_name()
    else:
        return save_name

def truncate_name(file_name):
    # Might not be unique, but looks better; user specifies the order anyway.
    path_seps = r'([/\\])'
    file_name = file_name.rsplit('.', 1)[0]
    match = re.search( path_seps , file_name )
    if match:
        file_name = '/'.join(file_name.rsplit(match.group(), 2)[-2:])
    return file_name

def main(myoptions):
    artists = []
    if not matplot:
        save = {}
    plotter = PlotRarefaction(myoptions)

    # Draw pre-loaded results and exit.
    if myoptions['--draw']:
        all_results = {}
        for saved_file in myoptions['-in']:
            with file(saved_file, 'rb') as infile:
                results = pickle.load(infile)
            all_results.update(results['results'])
        oldoptions = results['options']

        print(myoptions['--indices'], myoptions['-indices'])
        print(oldoptions['--indices'], oldoptions['-indices'])
        # Any other options need carrying over??
        oldoptions.update({ '-indices' : myoptions['-indices'],
                            '--indices': myoptions['--indices'] })
        plotter.options = oldoptions
        plotter.resume(all_results)
        return

    separate_processes = {}
    semaphore = multiprocessing.Semaphore(multiprocessing.cpu_count())
    update = separate_processes.update

    # Load data. Don't start threads yet. Got some counting to
    # do first.
    for infile in myoptions['-in']:
        with file( infile, 'rb' ) as in_handle:
            data_dict = pickle.load(in_handle)

        data_name = truncate_name(infile)

        print('Loading {0}'.format(data_name))
        separate_queue = multiprocessing.Queue()

        plotter.dataset_names.append(data_name)
        separate_process = Loader(data_dict, myoptions,
                queue=separate_queue, semaphore=semaphore)

        update({data_name : (separate_process, separate_queue, )})

    if myoptions['-increment'] is None:
        print('Pre-analysing result files to determine a universal increment (go for 20 steps).')
        increment = calc_increment(proc[0].nSeqs for proc in \
                                   separate_processes.values())
        plotter.increment = increment
        print('Sampling every {0} sequences.'.format(increment))

    # start the processes!!
    for process in separate_processes.itervalues():
        semaphore.acquire()
        process[0].increment = increment
        process[0].start()

    # Collate results and close processes.
    n_done = 0
    save = { 'options' : myoptions.options ,
             'results' : {} }
    processes_left = set(separate_processes.keys())
    while n_done < len(separate_processes):
        for data_name in list(processes_left):
            process, queue = separate_processes[data_name]
            try:
                data = queue.get(timeout=3)
                print('Rarefied {0}.'.format(data_name))
                n_done += 1
            except (Queue.Empty, IOError):
                continue
            processes_left.remove(data_name)
            data.file_name = data_name
            process.join()
            if matplot:
                plotter.load_rarefaction(data)
            save['results'].update({ data_name : data })

    # If we can, plot.
    if matplot:
        plotter.draw()

    # Save the results to plot later if you want...
    if myoptions['-out'] == []:
        myoptions['-out'] = [get_save_name()]
    if myoptions['-out'] != ([] or [None]):
        with file(os.path.realpath(myoptions['-out'][0]), 'wb') as outfile:
            pickle.dump( save , outfile , -1 )

if __name__ == '__main__':
    args = sys.argv[1:]
    myoptions = LocalOptions(args)
    main(myoptions)
else:
    myoptions = LocalOptions()
